import cv2
import numpy as np
import glob
import os

def crop_black(pano):
    """自动裁剪黑边"""
    gray = cv2.cvtColor(pano, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return pano
    x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
    return pano[y:y+h, x:x+w]


def stitch_images(path, out_name="stitched_optimized.jpg", mode="scans"):
    # 读取目录下所有 jpg 图片
    images = sorted(glob.glob(os.path.join(path, "*.jpg")))
    if len(images) < 2:
        print("❌ 目录中图像不足 2 张！")
        return

    print("检测到 {} 张图像".format(len(images)))
    imgs = [cv2.imread(p) for p in images]

    # 选择模式
    if mode.lower() == "panorama":
        stitch_mode = cv2.Stitcher_PANORAMA
    else:
        stitch_mode = cv2.Stitcher_SCANS   # 无人机推荐用 SCANS

    try:
        stitcher = cv2.Stitcher_create(stitch_mode)
    except:
        stitcher = cv2.createStitcher(stitch_mode)

    # 执行拼接
    status, pano = stitcher.stitch(imgs)

    if status == cv2.STITCHER_OK:
        # 自动裁剪黑边
        pano_cropped = crop_black(pano)

        out_path = os.path.join(path, out_name)
        cv2.imwrite(out_path, pano_cropped)
        print("✅ 优化拼接完成，保存至:", out_path)

        cv2.imshow("Optimized Panorama", pano_cropped)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("❌ 拼接失败，错误码:", status)


if __name__ == "__main__":
    # 修改成你的图像路径
    img_path = "/home/chengjinlai/qiqi-guaiguai-xiaowanyi/opencv-图像拼接/jpg"
    # mode="scans" 适合无人机航拍，mode="panorama" 适合相机旋转
    stitch_images(img_path, mode="scans")

